import os
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import cv2
from skimage.color import rgb2lab, lab2rgb

class ColorizationDataset(Dataset):
    """灰度图像着色数据集"""
    def __init__(self, root_dir, transform=None, mode='train'):
        """
        参数:
            root_dir (str): 数据集根目录
            transform (callable, optional): 数据转换
            mode (str): 'train' 或 'val'
        """
        self.root_dir = root_dir
        self.transform = transform
        self.mode = mode
        
        # 获取所有图像文件路径
        self.image_paths = []
        for root, _, files in os.walk(os.path.join(root_dir, mode)):
            for file in files:
                # 只加载彩色图像（不包含_gray后缀的图像）
                if file.lower().endswith(('.png', '.jpg', '.jpeg')) and not file.lower().endswith('_gray.png') and not file.lower().endswith('_gray.jpg') and not file.lower().endswith('_gray.jpeg'):
                    self.image_paths.append(os.path.join(root, file))
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        # 获取彩色图像路径和对应的灰度图像路径
        color_img_path = self.image_paths[idx]
        gray_img_path = color_img_path.replace('.jpg', '_gray.jpg').replace('.png', '_gray.png').replace('.jpeg', '_gray.jpeg')
        
        # 读取彩色图像
        color_img = Image.open(color_img_path).convert('RGB')
        
        # 应用变换
        if self.transform:
            color_img = self.transform(color_img)
        
        # 转换为LAB颜色空间
        img_lab = rgb2lab(color_img.permute(1, 2, 0).numpy())
        
        # 分离L通道和ab通道
        L = img_lab[:, :, 0] / 50.0 - 1.0  # 归一化到[-1, 1]
        ab = img_lab[:, :, 1:] / 110.0  # 归一化到[-1, 1]
        
        # 转换为张量
        L = torch.from_numpy(L).unsqueeze(0).float()
        ab = torch.from_numpy(ab).permute(2, 0, 1).float()
        
        return {'L': L, 'ab': ab, 'path': color_img_path}

def get_data_loaders(data_dir, batch_size=8, img_size=256, num_workers=4):
    """
    创建训练和验证数据加载器
    
    参数:
        data_dir (str): 数据集目录
        batch_size (int): 批次大小
        img_size (int): 图像大小
        num_workers (int): 数据加载线程数
    
    返回:
        train_loader, val_loader: 训练和验证数据加载器
    """
    # 数据转换
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
    ])
    
    # 创建数据集
    train_dataset = ColorizationDataset(data_dir, transform=transform, mode='train')
    val_dataset = ColorizationDataset(data_dir, transform=transform, mode='val')
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
    )
    
    return train_loader, val_loader

def lab_to_rgb(L, ab):
    """
    将LAB颜色空间的L和ab通道转换为RGB图像
    
    参数:
        L (torch.Tensor): L通道，形状为 [batch_size, 1, H, W]
        ab (torch.Tensor): ab通道，形状为 [batch_size, 2, H, W]
    
    返回:
        rgb_imgs (numpy.ndarray): RGB图像，形状为 [batch_size, H, W, 3]
    """
    L = (L + 1.0) * 50.0
    ab = ab * 110.0
    
    L = L.cpu().numpy()
    ab = ab.cpu().numpy()
    
    rgb_imgs = []
    for i in range(L.shape[0]):
        L_sample = L[i][0]
        ab_sample = np.transpose(ab[i], (1, 2, 0))
        
        # 合并L和ab通道
        lab_sample = np.concatenate([L_sample[:, :, np.newaxis], ab_sample], axis=2)
        
        # 转换为RGB
        rgb_sample = lab2rgb(lab_sample)
        rgb_imgs.append(rgb_sample)
    
    return np.array(rgb_imgs)

def save_checkpoint(model, optimizer, epoch, filename):
    """
    保存模型检查点
    
    参数:
        model: 模型
        optimizer: 优化器
        epoch: 当前轮次
        filename: 保存文件名
    """
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }
    torch.save(checkpoint, filename)

def load_checkpoint(filename, model, optimizer):
    """
    加载模型检查点
    
    参数:
        filename: 加载文件名
        model: 模型
        optimizer: 优化器
    
    返回:
        epoch: 加载的轮次
    """
    checkpoint = torch.load(filename)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    return checkpoint['epoch']

def prepare_grayscale_image(image_path, img_size=256):
    """
    准备灰度图像用于着色
    
    参数:
        image_path (str): 图像路径
        img_size (int): 图像大小
    
    返回:
        L (torch.Tensor): L通道，形状为 [1, 1, H, W]
    """
    # 读取图像
    img = Image.open(image_path).convert('RGB')
    
    # 调整大小
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
    ])
    img = transform(img)
    
    # 转换为LAB颜色空间
    img_lab = rgb2lab(img.permute(1, 2, 0).numpy())
    
    # 提取L通道
    L = img_lab[:, :, 0] / 50.0 - 1.0  # 归一化到[-1, 1]
    
    # 转换为张量
    L = torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float()
    
    return L 

def prepare_and_restore_grayscale_image(image_path, model_img_size=256):
    """
    准备灰度图像用于着色，并返回原始尺寸信息用于后续还原
    
    参数:
        image_path (str): 图像路径
        model_img_size (int): 模型处理的图像大小
    
    返回:
        L (torch.Tensor): L通道，形状为 [1, 1, H, W]
        original_size (tuple): 原始图像尺寸 (width, height)
        original_image (PIL.Image): 原始图像对象
    """
    # 读取图像
    original_image = Image.open(image_path).convert('RGB')
    original_size = original_image.size  # (width, height)
    
    # 调整大小用于模型处理
    transform = transforms.Compose([
        transforms.Resize((model_img_size, model_img_size)),
        transforms.ToTensor(),
    ])
    img = transform(original_image)
    
    # 转换为LAB颜色空间
    img_lab = rgb2lab(img.permute(1, 2, 0).numpy())
    
    # 提取L通道
    L = img_lab[:, :, 0] / 50.0 - 1.0  # 归一化到[-1, 1]
    
    # 转换为张量
    L = torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float()
    
    return L, original_size, original_image

def lab_to_rgb_with_resize(L, ab, original_size):
    """
    将LAB颜色空间的L和ab通道转换为RGB图像，并调整为原始大小
    
    参数:
        L (torch.Tensor): L通道，形状为 [batch_size, 1, H, W]
        ab (torch.Tensor): ab通道，形状为 [batch_size, 2, H, W]
        original_size (tuple): 原始图像尺寸 (width, height)
    
    返回:
        rgb_imgs (numpy.ndarray): 调整为原始大小的RGB图像，形状为 [batch_size, H, W, 3]
    """
    # 先转换为RGB
    rgb_imgs = lab_to_rgb(L, ab)
    
    # 调整大小
    resized_imgs = []
    for img in rgb_imgs:
        # 将浮点数值限制在[0,1]范围内
        img = np.clip(img, 0, 1)
        # 转换为PIL图像
        pil_img = Image.fromarray((img * 255).astype(np.uint8))
        # 调整为原始大小
        pil_img = pil_img.resize(original_size, Image.BILINEAR)
        # 转回numpy数组并归一化到[0,1]
        resized_img = np.array(pil_img).astype(np.float32) / 255.0
        resized_imgs.append(resized_img)
    
    return np.array(resized_imgs) 